import java.security.Key;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.KeyFactory;
import java.security.Signature;
import java.security.GeneralSecurityException;
import java.security.spec.RSAPublicKeySpec;
import java.security.spec.RSAPrivateKeySpec;
import javax.crypto.Cipher;
import javax.crypto.KeyGenerator;
import javax.crypto.spec.IvParameterSpec;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.ObjectInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.math.BigInteger;

/**
 *
 @author Paul Ozog
 Encrypt/Decrypts a file

 -Encryption:
 Uses a destination host's public RSA key to
 encrypt an AES key.  The AES key in turn is used to encrypt 
 a plaintext file.  Finally, [wrapped AES key, ciphertext] bytes are hashsed
 and signed with the sender's private key.  All three sets of 
 bits are placed in the output file

 -Decryption:
 Reads the three sets of bits described above. Decrypts the AES key,
 which decrypts the ciphertext.  Finally, the signature is verified
 by checking the SHA1 sum of the [wrapped AES key, ciphertext] with
 the sender's public key.  If the result is the same as in the
 Encryption process, then the signature is valid.  If not, then
 the program says so, exits, and the output plaintext file is not created.

 -Remarks on encryption modes, keysizes:
 RSA uses the standard 1024 key length, and is generated by the java
 KeyPairGenerator.  Keys are to be made with the included genKeys.java + Makefile
 AES uses the default key size
 CBC mode is used instead of the simpler but less secure ECB. 
 Signature verification is done with SHA1WithRSA - we hash the entire 
 AESKey + ciphertext and use that as our signature
*/ 

class fcrypt {
    //Static members for both encryption and decryption
    private static Cipher AESCipher, RSACipher;
    private static Key AESKey;
    private static byte[] ciphertext;
    private static byte[] plaintext;
    private static byte[] IV;

    //Static members for encryption:
    private static KeyGenerator AESKeyGen;
    private static PublicKey RSAPublicKey;

    //Static members for decryption:
    private static PrivateKey RSAPrivateKey;

    /**
     @param args - see printUsage()
     @return none

     Check args and call either beginEncrypt or beginDecrypt, depending on
     argv[0]
     
    */
    public static void main (String[] args) {
	if (args.length != 5) {
	    printUsage();
	    System.exit(0);
	}
	try {
	    if (args[0].equals("-e"))
		beginEncrypt(args);
	    else if (args[0].equals("-d"))
		beginDecrypt(args);
	    else
		printUsage();
	} catch (IndexOutOfBoundsException e) {
	    printUsage();
	}
	finally {
	    System.exit(1);
	}

    } // end main

    /**
       @param None
       @return None
       
       Print the usage of fcrypt

    */
    public static void printUsage() {
	System.out.println("Usage:");
	System.out.println("java fcrypt -e destination_public_key_filename" + 
			   " sender_private_key_filename input_plaintext_file" +
			   " output_ciphertext_file\n");
	System.out.println("-------or-------");
	System.out.println("\njava fcrypt -d destination_private_key_filename" + 
			   " sender_public_key_filename input_ciphertext_file" +
			   " output_plaintext_file");
    }

    /**
       @param args - arguments to fcrypt main
       @return None

       Calls the major pieces of the encryption process
       
    */
    private static void beginEncrypt(String[] args) {
	String destPublicKey      = args[1];
	String senderPrivateKey   = args[2];
	String plaintextFileName  = args[3];
	String ciphertextFileName = args[4];

	try {
	    createCiphers();
	    createAESKey();
	    RSAPublicKey = readPubKeyFromFile(new File(destPublicKey));
	    RSAPrivateKey= readPrivKeyFromFile(new File(senderPrivateKey));
	    plaintext  = getFileBytes(new File(plaintextFileName));
	    ciphertext = AESEncrypt(plaintext);
	    dumpToFile(ciphertextFileName);
	} catch (Exception e) {
	    e.printStackTrace();
	    System.exit(1);
	}

    } //end beginEncrypt

    /**
       @param args - arguments to fcrypt main
       @return None

       Calls the major pieces of the decryption process
       
    */
    private static void beginDecrypt(String[] args) {
	String destPrivateKey     = args[1];
	String senderPublicKey    = args[2];
	String ciphertextFileName = args[3];
	String plaintextFileName  = args[4];

	try {
	    createCiphers();
	    RSAPublicKey  = readPubKeyFromFile(new File(senderPublicKey));
	    RSAPrivateKey = readPrivKeyFromFile(new File(destPrivateKey));
	    plaintext = readEncryptedFile(new File(ciphertextFileName));
	    writeToFile(plaintext, new File(plaintextFileName));
	}  catch (Exception e) {
	    e.printStackTrace();
	    System.exit(1);
	}

    } // end beginDecrypt
	
    /**
       @param None
       @return None
       @throws GeneralSecurityException
       Creates two ciphers - one for AES and the other for RSA
       
    */
    private static void createCiphers() throws GeneralSecurityException {
	AESCipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
	RSACipher = Cipher.getInstance("RSA");
    }

    /**
       @param None
       @return None
       @throws GeneralSecurityException
       
       Generates the AESKey using a KeyGenerator

    */
    private static void createAESKey() throws GeneralSecurityException {
	AESKeyGen = KeyGenerator.getInstance("AES");
	AESKey = AESKeyGen.generateKey();
    }
    
    /**
       @param keyFile - File that contains key made from KeyPairGenerator
       @return PublicKey
       @throws IOException

       Based on example from Neil Coffey's "RSA Encryption in Java"
       (www.javamex.com). Also see genKeys.java

    */
    private static PublicKey readPubKeyFromFile(File keyFile) throws IOException {
	FileInputStream fis = new FileInputStream(keyFile);
	ObjectInputStream oin = new ObjectInputStream(fis);

	try {
	    BigInteger m = (BigInteger) oin.readObject();
	    BigInteger e = (BigInteger) oin.readObject();
	    RSAPublicKeySpec keySpec = new RSAPublicKeySpec(m, e);
	    KeyFactory fact = KeyFactory.getInstance("RSA");
	    PublicKey pubKey = fact.generatePublic(keySpec);
	    return pubKey;
	} catch (Exception e) {
	    throw new RuntimeException("Spurious serialisation error", e);
	} finally {
	    oin.close();
	}

    } // end readPubKeyFromFile

    /**
       @param keyFile - File that contains key made from KeyPairGenerator
       @return PrivateKey
       @throws IOException

       Similar to readPubKeyFromFile, except read a Private key. 
       Also see genKeys.java

    */
    private static PrivateKey readPrivKeyFromFile(File keyFile) throws IOException {
	FileInputStream fis = new FileInputStream(keyFile);
	ObjectInputStream oin = new ObjectInputStream(fis);

	try {
	    BigInteger m = (BigInteger) oin.readObject();
	    BigInteger e = (BigInteger) oin.readObject();
	    RSAPrivateKeySpec keySpec = new RSAPrivateKeySpec(m, e);
	    KeyFactory fact = KeyFactory.getInstance("RSA");
	    PrivateKey privKey = fact.generatePrivate(keySpec);
	    return privKey;
	} catch (Exception e) {
	    throw new RuntimeException("Spurious serialisation error", e);
	} finally {
	    oin.close();
	}

    } // end readPrivKeyFromFile

    /**
       @param bytes - plaintext bytes to encrypt
       @return ciphertext - encrypted bytes (byte[])
       @throws GeneralSecurityException

       Uses the AESCipher to encrypt the input bytes.  Returns the ciphertext
       bytes
       
    */
    private static byte[] AESEncrypt (byte[] bytes) throws GeneralSecurityException {
	// Initialize the cipher for encryption
	AESCipher.init(Cipher.ENCRYPT_MODE, AESKey);
	
	//get an initialization vector:
	IV = AESCipher.getIV();

	// Encrypt the cleartext
	ciphertext = AESCipher.doFinal(bytes);
	return ciphertext;
    }

    /**
       @param bytes - ciphertext bytes to encrypt
       @return plaintext - decrypted bytes (byte[])
       @throws GeneralSecurityException

       Uses the AESCipher to decrypt the input bytes.  Returns 
       plaintext

       This functions assumes that IV is initiated with a call
       to readEncryptedFile
       
    */
    private static byte[] AESDecrypt (byte[] bytes) throws GeneralSecurityException {
	byte[] message;
	// Initialize the cipher with the key and IV for decryption
	AESCipher.init(Cipher.DECRYPT_MODE, AESKey, new IvParameterSpec(IV));

	// Get the decrypted plaintext
	message = AESCipher.doFinal(bytes);

	return message;
    }

    /**
       @param key - Key to wrap using RSA
       @return wrapped key - byte[]
       @throws GeneralSecurityException

       This should only be called for wrapping the AESKey
       
    */
    private static byte[] RSAWrap (Key key) throws GeneralSecurityException {
	byte[] wrappedKey;
	// Initialize the cipher for wrapping
	RSACipher.init(Cipher.WRAP_MODE, RSAPublicKey);

	// Wrap the key
	wrappedKey = RSACipher.wrap(key);

	return wrappedKey;
    }

    /**
       @param bytes - byte[] of wrapped Key
       @return key  - Key that represents unwrapped bytes
       @throws GeneralSecurityException

       This should only be called against wrapped AESKey bytes

    */
    private static Key RSAUnwrap (byte[] bytes) throws GeneralSecurityException {
	Key key;
	// Initialize the cipher for unwrapping
	RSACipher.init(Cipher.UNWRAP_MODE, RSAPrivateKey);

	// Wrap the bytes and return the key
	key = RSACipher.unwrap(bytes, "AES", Cipher.SECRET_KEY);

	return key;
    }

    /**
       @param file - File to convert to byte array
       @return fileBytes - byte[] of file's contents
       @throws IOException
      
       Pretty straightforward - convert File to byte[]

    */
    private static byte[] getFileBytes(File file) throws IOException {
	byte[] fileBytes = new byte[(int)file.length()];
	
	try {
	    FileInputStream fis = new FileInputStream(file);
	    fis.read(fileBytes);
	    fis.close();
	} catch (IOException e) {
	    e.printStackTrace();
	    System.exit(1);
	}

	return fileBytes;
    } // end getFileBytes

    /**
       @param fileName   - string to dump ciphertext
       @return None
      
       Write the necessary data fields into an encrypted file
       Encrypted file format will have the following form: 
	   [AESKey length] [IV length] [ciphertext length] \
	   [sigBytes length] [AESKey] [IV] [ciphertext] [sigBytes]

    */
    private static void dumpToFile (String fileName) {

	try {
	    FileOutputStream fos = new FileOutputStream(fileName);
	    DataOutputStream dos = new DataOutputStream(fos);

	    byte[] AESKeyEncrypted = RSAWrap(AESKey);

	    Signature sig = Signature.getInstance("SHA1WithRSA");
	    sig.initSign(RSAPrivateKey);
	    /* Update the signature with the unencrypted AES key
	       and plaintext bytes concatenated together*/
	    sig.update(concat(AESKeyEncrypted, ciphertext));

	    byte[] sigBytes = sig.sign();

	    dos.writeInt(AESKeyEncrypted.length);
	    dos.writeInt(IV.length);
	    dos.writeInt(ciphertext.length);
	    dos.writeInt(sigBytes.length);
	    dos.write(AESKeyEncrypted);
	    dos.write(IV);
	    dos.write(ciphertext);
	    dos.write(sigBytes);

	} catch (Exception e) {
	    e.printStackTrace();
	    System.exit(1);
	} // end try catch

    } // end dumpToFile

    /**
       @param file - encrypted File to read
       @return plaintext - byte[]
       @throws IOException

       Basically the inverse of dumpToFile

       Also, check the signature.  If not valid, exit the program.  Else,
       keep executing
    */
    private static byte[] readEncryptedFile(File file) throws IOException {
	byte[] AESKeyRawEncrypted, sigData;
	byte[] message = new byte[0];
	try {
	    FileInputStream fis = new FileInputStream(file);
	    DataInputStream dis = new DataInputStream(fis);
	    Signature sig;

	    AESKeyRawEncrypted = new byte[dis.readInt()];
	    IV = new byte[dis.readInt()];
	    ciphertext= new byte[dis.readInt()];
	    sigData = new byte[dis.readInt()];
	    
	    dis.read(AESKeyRawEncrypted);
	    dis.read(IV);
	    dis.read(ciphertext);
	    dis.read(sigData);

	    sig = Signature.getInstance("SHA1WithRSA");
	    sig.initVerify(RSAPublicKey);
	    /* Update the signature with the AES key
	       and plaintext bytes concatenated together*/
	    sig.update(concat(AESKeyRawEncrypted, ciphertext));

	    if (sig.verify(sigData))
	    	System.out.println("The signature is valid");
	    else {
	    	System.out.println("The signature is NOT valid! Exiting...");
	    	System.exit(-1);
	    }

	    AESKey = RSAUnwrap(AESKeyRawEncrypted);
	    
	    message = AESDecrypt(ciphertext);

	} catch (Exception e) {
	    e.printStackTrace();
	    System.exit(1);
	} // end try catch

	return message;

    } // end readEncryptedFile

    
    /**
       @param data - byte[] that will be written to fileName file
       @param fileName - String of the file to write s
       @throws Exception
    */
    private static void writeToFile(byte[] data, File f) throws Exception {
	FileOutputStream fos = new FileOutputStream(f);
	DataOutputStream dos = new DataOutputStream(fos);
	dos.write(data);
	dos.close();
    }

    /**
       @param a - byte[] 
       @param b - byte[]
       @return result - byte[] of a and b concatenated together

       Used to concat the encrypted AES key and ciphertext, to feed
       into Signature verification
       
    */
    private static byte[] concat(byte[] a, byte[] b) {
	final int alen = a.length;
	final int blen = b.length;
	if (alen == 0) {
	    return b;
	}
	if (blen == 0) {
	    return a;
	}
	final byte[] result = (byte[]) java.lang.reflect.Array.
            newInstance(a.getClass().getComponentType(), alen + blen);
	System.arraycopy(a, 0, result, 0, alen);
	System.arraycopy(b, 0, result, alen, blen);
	return result;
    } // end concat

} // end fcrypt
